package pl.edu.agh.anomalyids.lerad.rule;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Random;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import pl.edu.agh.anomalyids.lerad.aspect.AbstractAspect;
import pl.edu.agh.anomalyids.lerad.aspect.enums.Match;
import pl.edu.agh.anomalyids.lerad.aspect.implementation.CodeIcmpAspect;
import pl.edu.agh.anomalyids.lerad.aspect.implementation.DestIpAddressAspect;
import pl.edu.agh.anomalyids.lerad.aspect.implementation.DestPortAspect;
import pl.edu.agh.anomalyids.lerad.aspect.implementation.HttpMethodAspect;
import pl.edu.agh.anomalyids.lerad.aspect.implementation.SrcIpAddressAspect;
import pl.edu.agh.anomalyids.lerad.aspect.implementation.SrcPortAspect;
import pl.edu.agh.anomalyids.lerad.aspect.implementation.TcpFlagAspect;
import pl.edu.agh.anomalyids.lerad.aspect.implementation.TtlAspect;
import pl.edu.agh.anomalyids.lerad.aspect.implementation.TypeIcmpAspect;
import pl.edu.agh.anomalyids.lerad.rule.util.PacketsDivider;
import pl.edu.agh.anomalyids.lerad.rule.util.RuleBeliefComparator;
import pl.edu.agh.anomalyids.lerad.util.ImplicationValueAdder;
import pl.edu.agh.anomalyids.network.Packet;
import pl.edu.agh.anomalyids.network.Packet.Protocol;
import pl.edu.agh.anomalyids.network.header.impl.HTTPHeader;
import pl.edu.agh.anomalyids.network.header.impl.ICMPHeader;
import pl.edu.agh.anomalyids.network.header.impl.IPv4Header;
import pl.edu.agh.anomalyids.network.header.impl.TCPHeader;
import pl.edu.agh.anomalyids.network.header.impl.UDPHeader;

public class RuleCreator {
	
	private static final double REWARD = 0.5;
	private static final double PENALTY = 0.5;

	private final List<Rule> rules;
	private Map<Integer, List<Rule>> rulesMap;
	private PacketsDivider packetDivider;
	private Logger logger = LoggerFactory.getLogger(RuleCreator.class);

	public RuleCreator(List<Packet> packets) {
		this.rules = new ArrayList<Rule>();
		this.rulesMap = new HashMap<Integer, List<Rule>>();
		this.packetDivider = new PacketsDivider(packets);
	}
	
	public RuleCreator generateRules() {
		logger.info("Generating rules for sample set of packets...");
		List<Packet> packetsList = packetDivider.getSampleSetList();
		
		Random random = new Random();
		
		int packetsSize = packetsList.size();
		int randomNumber1;
		int randomNumber2;
		for (int iteration = 0; iteration < packetsSize; ++iteration) {
			randomNumber1 = random.nextInt(packetsSize);
			randomNumber2 = random.nextInt(packetsSize);
			
			Rule rule = createRule(packetsList.get(randomNumber1), packetsList.get(randomNumber2), packetsList);
			if (rule != null) {
				addRule(rule);
			}
		}
		
		Collections.sort(rules, new RuleBeliefComparator());
		
		logger.info("DONE");
		
		return this;
	}
	
	public RuleCreator performCoverage() {
		logger.info("Performing coverage test...");
		
		List<Packet> packetsList = packetDivider.getSampleSetList();
		
		Map<Integer, Integer> coverageMap = new HashMap<Integer, Integer>();
		
		int ruleSize = rules.size();
		for (Packet packet : packetsList) {
			for (int ruleIterator = 0; ruleIterator < ruleSize; ++ruleIterator) {
				Collection<AbstractAspect<?>> aspects = rules.get(ruleIterator).getConditionals();
				
				boolean hasPassed = doesLeftSideMatch(aspects, packet);
				
				if (hasPassed) {
					hasPassed = doesRightSideMatch(rules.get(ruleIterator).getImplication(), packet);
					
					if (hasPassed) {
						updateCoverageMap(ruleIterator, coverageMap);
						break;
					}
				}
			}
		}
		
		removeRedundantRules(coverageMap);
		
		logger.info("DONE");
		
		return this;
	}
	
	public RuleCreator train() {
		logger.info("Training rules on training set of packets...");
		
		List<Packet> packetsList = packetDivider.getTrainingSetList();
		
		int ruleSize = rules.size();
		for (Packet packet : packetsList) {
			for (int ruleIterator = 0; ruleIterator < ruleSize; ++ruleIterator) {
				Collection<AbstractAspect<?>> aspects = rules.get(ruleIterator).getConditionals();
				
				boolean hasPassed = doesLeftSideMatch(aspects, packet);
				
				if (hasPassed) {
					Rule rule = rules.get(ruleIterator);
					
					rule.setNValue(rule.getNValue() + 1);
					
					ImplicationValueAdder.tryAdd(packet, rule.getImplication());
				}
			}
		}
		
		Collections.sort(rules, new RuleBeliefComparator());
		
		logger.info("DONE");
		
		return this;
	}
	
	private void addRule(Rule rule) {
		List<Rule> currentRules = rulesMap.get(rule.hashCode());
		
		if (currentRules != null) {
			for (Rule mapRule : currentRules) {
				if (mapRule.getImplication().getType() == rule.getImplication().getType()) {
					if (!mapRule.getImplication().getValues().contains(rule.getImplication().getData())) {
						Collection collection = mapRule.getImplication().getValues();
						collection.add(rule.getImplication().getData());
						mapRule.setNValue(mapRule.getNValue() + rule.getNValue());
					}
					
					return;
				}
			}
			
			rules.add(rule);
			currentRules.add(rule);
		} else {
			List<Rule> newRulesList = new ArrayList<Rule>();
			newRulesList.add(rule);
			
			rules.add(rule);
			rulesMap.put(rule.hashCode(), newRulesList);
		}
	}
	
	private void removeRedundantRules(Map<Integer, Integer> coverageMap) {
		List<Rule> rulesToRemove = new ArrayList<Rule>();
		
		int rulesSize = rules.size();
		for (int rulesIterator = 0; rulesIterator < rulesSize; ++ rulesIterator) {
			if (!coverageMap.containsKey(rulesIterator)) {
				rulesToRemove.add(rules.get(rulesIterator));
			}
		}
		
		for (Rule rule : rulesToRemove) {
			rules.remove(rule);
		}
	}
	
	private void updateCoverageMap(int rule, Map<Integer, Integer> coverageMap) {
		Integer value = coverageMap.get(rule);
		
		if (value != null) {
			coverageMap.put(rule, value.intValue() + 1);
		} else {
			coverageMap.put(rule, 1);
		}
	}
	
	private boolean doesLeftSideMatch(Collection<AbstractAspect<?>> aspects, Packet packet) {
		boolean hasPassed = true;
		
		for (AbstractAspect<?> aspect : aspects) {
			if (aspect.packetMatchesAspect(packet) == Match.DIFFER) {
				hasPassed = false;
				break;
			}
		}
		
		return hasPassed;
	}
	
	private boolean doesRightSideMatch(AbstractAspect<?> aspect, Packet packet) {
		return aspect.packetMatchesAspect(packet) == Match.SAME ? true : false;
	}

	public Rule createRule(Packet packet1, Packet packet2, List<Packet> packetsList) {
		List<AbstractAspect<?>> aspects = new ArrayList<AbstractAspect<?>>(); 
		for (Protocol protocol : packet1.getProtocols().keySet()) {
			if (packet2.getProtocols().containsKey(protocol)) {
				List<AbstractAspect<?>> aspectsList = createAspects(packet1, packet2, protocol);
				
				aspects.addAll(aspectsList);
			}
		}
		
		if (aspects.size() > 1) {
			Rule rule = createRuleWithAspects(aspects);
			rule.setNValue(getNValue(rule.getConditionals(), packetsList));
			
			return rule;
		}
		
		return null;
	}
	
	private int getNValue(Collection<AbstractAspect<?>> aspects, List<Packet> packetsList) {
		int packetPassed = 0;
		
		for (Packet packet : packetsList) {
			boolean hasPassed = doesLeftSideMatch(aspects, packet);
			
			if (hasPassed) {
				++packetPassed;
			}
		}
		
		return packetPassed;
	}
	
	private Rule createRuleWithAspects(List<AbstractAspect<?>> aspects) {
		Rule rule = new Rule();
		
		rule.setImplication(getAndRemoveAspect(aspects, new Random().nextInt(aspects.size())));
		
		for (AbstractAspect<?> aspect : aspects) {
			rule.addAspect(aspect);
		}
		
		return rule;
	}
	
	private List<AbstractAspect<?>> createAspects(Packet packet1, Packet packet2, Protocol protocol) {	
		List<AbstractAspect<?>> aspects = new ArrayList<AbstractAspect<?>>();
		
		if (protocol == Protocol.IPv4) {
			IPv4Header header1 = (IPv4Header) packet1.getProtocols().get(protocol);
			IPv4Header header2 = (IPv4Header) packet2.getProtocols().get(protocol);
			
			addIPv4Aspects(aspects, header1, header2);
		} else if (protocol == Protocol.HTTP) {
			HTTPHeader header1 = (HTTPHeader) packet1.getProtocols().get(protocol);
			HTTPHeader header2 = (HTTPHeader) packet2.getProtocols().get(protocol);
			
			addHTTPAspects(aspects, header1, header2);
		} else if (protocol == Protocol.ICMP) {
			ICMPHeader header1 = (ICMPHeader) packet1.getProtocols().get(protocol);
			ICMPHeader header2 = (ICMPHeader) packet2.getProtocols().get(protocol);
			
			addICMPAspects(aspects, header1, header2);
		} else if (protocol == Protocol.TCP) {
			TCPHeader header1 = (TCPHeader) packet1.getProtocols().get(protocol);
			TCPHeader header2 = (TCPHeader) packet2.getProtocols().get(protocol);
			
			addTCPAspects(aspects, header1, header2);
		} else if (protocol == Protocol.UDP) {
			UDPHeader header1 = (UDPHeader) packet1.getProtocols().get(protocol);
			UDPHeader header2 = (UDPHeader) packet2.getProtocols().get(protocol);
			
			addUDPAspects(aspects, header1, header2);
		} 
		
		return aspects;
	}
	
	private void addIPv4Aspects(List<AbstractAspect<?>> aspects, IPv4Header header1, IPv4Header header2) {
		if (header1.getSrcAddress().equals(header2.getSrcAddress())) {
			aspects.add(new SrcIpAddressAspect(header1.getSrcAddress()));
		}
		
		if (header1.getDestAddress().equals(header2.getDestAddress())) {
			aspects.add(new DestIpAddressAspect(header1.getDestAddress()));
		}
		
		if (header1.getTTL() == header2.getTTL()) {
			aspects.add(new TtlAspect(header1.getTTL()));
		}
	}
	
	private void addHTTPAspects(List<AbstractAspect<?>> aspects, HTTPHeader header1, HTTPHeader header2) {
		if (header1.getRequestMethod().equals(header2.getRequestMethod())) {
			aspects.add(new HttpMethodAspect(header1.getRequestMethod()));
		}
	}
	
	private void addTCPAspects(List<AbstractAspect<?>> aspects, TCPHeader header1, TCPHeader header2) {
		if (header1.getDestPort() == header2.getDestPort() ) {
			aspects.add(new DestPortAspect(header1.getDestPort()));
		}
		
		if (header1.getSrcPort() == header2.getSrcPort()) {
			aspects.add(new SrcPortAspect(header1.getSrcPort()));
		}
		
		if (header1.getFlags() == header2.getFlags()) {
			aspects.add(new TcpFlagAspect(header1.getFlags()));
		}
	}
	
	private void addUDPAspects(List<AbstractAspect<?>> aspects, UDPHeader header1, UDPHeader header2) {
		if (header1.getDestPort() == header2.getDestPort() ) {
			aspects.add(new DestPortAspect(header1.getDestPort()));
		}
		
		if (header1.getSrcPort() == header2.getSrcPort()) {
			aspects.add(new SrcPortAspect(header1.getSrcPort()));
		}
	}
	
	private void addICMPAspects(List<AbstractAspect<?>> aspects, ICMPHeader header1, ICMPHeader header2) {
		if (header1.getCode() == header2.getCode()) {
			aspects.add(new CodeIcmpAspect(header1.getCode()));
		}
		
		if (header1.getType() == header2.getType()) {
			aspects.add(new TypeIcmpAspect(header1.getType()));
		}
	}
	
	private AbstractAspect<?> getAndRemoveAspect(List<AbstractAspect<?>> aspects, int aspectNumber) {
		AbstractAspect<?> aspect = aspects.get(aspectNumber);
		aspects.remove(aspectNumber);
		
		return aspect;
	}

	public RuleCreator validate() {
		logger.info("Validating rules on validation set of packets...");
		
		List<Packet> packetsList = packetDivider.getValidationSetList();
		
		List<Rule> tempRulesList = new ArrayList<Rule>();
		boolean ruleViolated = false;
		
		for (Packet packet : packetsList){			
			for (Rule rule : rules){
				Match ruleMatch = rule.validate(packet);
				if (ruleMatch.equals(Match.SAME)){
					tempRulesList.add(rule);
				} else if (ruleMatch.equals(Match.DIFFER)){
					ruleViolated = true;
					rule.setWeight(rule.getWeight()*PENALTY);
				}
			}
			
			if (ruleViolated){
				for (Rule rule : tempRulesList){
					rule.setWeight(rule.getWeight()*(1.0 + REWARD));
				}
			}
			ruleViolated = false;
			tempRulesList.clear();
		}
		
		logger.info("DONE");
		
		return this;
	}

	public List<Rule> getRules() {
		return rules;
	}
}
